package spring.cloud.tasks.executor.listener.sharding;

import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import com.google.gson.Gson;
import com.google.gson.reflect.TypeToken;
import lombok.Getter;
import lombok.Setter;
import lombok.extern.slf4j.Slf4j;
import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.api.transaction.CuratorTransactionFinal;
import spring.cloud.tasks.executor.node.Node;

import java.nio.charset.StandardCharsets;
import java.util.*;

@Slf4j
@Setter
@Getter
public class AllotTaskItemContentService {

    public static final String GET_SHARD_FROM_CONTENT_NODE_FLAG = "GET_SHARD_FROM_CONTENT_NODE_FLAG";
    private CuratorFramework curatorFramework;

    private Gson gson = new Gson();

    public AllotTaskItemContentService(CuratorFramework curatorFramework) {
        this.curatorFramework = curatorFramework;
    }

    public void persistExecutorList(List<Executor> executorList) throws Exception {
        // sharding/content如果不存在，则新建
        if (curatorFramework.checkExists().forPath(Node.SHARDING_CONTENTNODE_PATH) == null) {
            curatorFramework.create().creatingParentsIfNeeded().forPath(Node.SHARDING_CONTENTNODE_PATH);
        }
        // 删除sharding/content节点下的内容
        List<String> childrenPathList = curatorFramework.getChildren().forPath(Node.SHARDING_CONTENTNODE_PATH);
        if (childrenPathList != null && !childrenPathList.isEmpty()) {
            for (String childrenPath : childrenPathList) {
                curatorFramework.delete().forPath(Node.getShardingContentElementNodePath(childrenPath));
            }
        }
        // 持久化新的内容
        List<List<Executor>> executorListList = Lists.partition(executorList, 16);
        // 如果内容过大，分开节点存储。不能使用事务提交，因为即使使用事务、写多个节点，但是提交事务时，仍然会报长度过长的错误。
        for (int i = 0; i < executorListList.size(); i++) {
            byte[] bytes = gson.toJson(executorListList.get(i)).getBytes(StandardCharsets.UTF_8);
            curatorFramework.create().forPath(Node.getShardingContentElementNodePath(String.valueOf(i)), bytes);
        }
    }


    /**
     * @param taskId 作业名
     * @return 返回Map数据，key值为executorName, value为分片项集合
     */
    public Map<String, List<Integer>> getExecutorIdToItemIndexListMap(String taskId) throws Exception {
        List<Executor> executorList = getExecutorList();
        return getExecutorIdToItemIndexListMap(executorList, taskId);
    }

    /**
     * 从sharding/content获取数据
     */
    public List<Executor> getExecutorList() throws Exception {
        List<Executor> executorList = new ArrayList<>();
        // Sharding/content 内容多的时候，分多个节点存数据
        if (curatorFramework.checkExists().forPath(Node.SHARDING_CONTENTNODE_PATH) != null) {
            List<String> childrenPathList = curatorFramework.getChildren().forPath(Node.SHARDING_CONTENTNODE_PATH);
            Collections.sort(childrenPathList, new Comparator<String>() {
                @Override
                public int compare(String arg0, String arg1) {
                    Integer a = Integer.valueOf(arg0);
                    Integer b = Integer.valueOf(arg1);
                    return a.compareTo(b);
                }
            });
            for (String childrenPath : childrenPathList) {
                byte[] bytes = curatorFramework.getData().forPath(Node.getShardingContentElementNodePath(childrenPath));
                executorList.addAll((Collection<? extends Executor>) gson.fromJson(new String(bytes, StandardCharsets.UTF_8), new TypeToken<List<Executor>>() {
                }.getType()));
            }
        }
        return executorList;
    }

    public Map<String, List<Integer>> getExecutorIdToItemIndexListMap(List<Executor> executorList, String taskId) {
        if (executorList == null || executorList.isEmpty()) {
            return Maps.newHashMap();
        }

        Map<String, List<Integer>> executorIdToItemIndexListMap = new HashMap<>();
        for (Executor executor : executorList) {
            if (executor.getTaskIdList() != null && executor.getTaskIdList().contains(taskId)) {
                List<Integer> taskItemList = new ArrayList<>();
                for (TaskItem taskItem : executor.getTaskItemList()) {
                    if (taskItem.getTaskId().equals(taskId)) {
                        taskItemList.add(taskItem.getItemIndex());
                    }
                }
                executorIdToItemIndexListMap.put(executor.getExecutorId(), taskItemList);
            }
        }
        return executorIdToItemIndexListMap;
    }

    /**
     * 事务
     *
     * @param taskIdToExecutorIdToItemIndexListMapMap
     * @throws Exception
     */
    public void persistTaskIdToExecutorIdToItemIndexListMapMap(Map<String, Map<String, List<Integer>>> taskIdToExecutorIdToItemIndexListMapMap) throws Exception {
        if (!taskIdToExecutorIdToItemIndexListMapMap.isEmpty()) {
            log.info("Notify jobs sharding necessary, jobs is {}", taskIdToExecutorIdToItemIndexListMapMap.keySet());
            CuratorTransactionFinal curatorTransactionFinal = curatorFramework.inTransaction().check().forPath("/").and();
            Iterator<Map.Entry<String, Map<String, List<Integer>>>> iterator = taskIdToExecutorIdToItemIndexListMapMap.entrySet().iterator();
            while (iterator.hasNext()) {
                Map.Entry<String, Map<String, List<Integer>>> next = iterator.next();
                String taskId = next.getKey();
                Map<String, List<Integer>> executorIdToItemIndexListMap = next.getValue();
                byte[] bytes = gson.toJson(executorIdToItemIndexListMap).getBytes(StandardCharsets.UTF_8);
                // 更新$Jobs/xx/leader/sharding/neccessary 节点的内容为新分配的sharding 内容
                String taskLeaderShardingNodePath = Node.getTaskLeaderShardingNodePath(taskId);
                if (curatorFramework.checkExists().forPath(taskLeaderShardingNodePath) == null) {
                    curatorFramework.create().creatingParentsIfNeeded().forPath(taskLeaderShardingNodePath);
                }
                String taskLeaderShardingNecessaryNodePath = Node.getTaskLeaderShardingNecessaryNodePath(taskId);
                if (curatorFramework.checkExists().forPath(taskLeaderShardingNecessaryNodePath) == null) {
                    curatorTransactionFinal.create().forPath(taskLeaderShardingNecessaryNodePath, bytes).and();
                } else {
                    curatorTransactionFinal.setData().forPath(taskLeaderShardingNecessaryNodePath, bytes).and();
                }
            }
            curatorTransactionFinal.commit();
        }
    }

    public Map<String, List<Integer>> getExecutorIdToItemIndexListMap(String taskId, String taskNecessaryContent) throws Exception {
        Map<String, List<Integer>> executorIdToItemIndexListMap = new HashMap<>();
        // data为'GET_SHARD_FROM_CONTENT_NODE_FLAG'，会直接从sharding/content节点获取数据
        if (GET_SHARD_FROM_CONTENT_NODE_FLAG.equals(taskNecessaryContent)) {
            executorIdToItemIndexListMap.putAll(getExecutorIdToItemIndexListMap(taskId));
            return executorIdToItemIndexListMap;
        }

        try {
            // TODO: can change to Jackson ?
            Map<String, List<Integer>> $executorIdToItemIndexListMap = gson.fromJson(taskNecessaryContent, new TypeToken<Map<String, List<Integer>>>() {
            }.getType());
            executorIdToItemIndexListMap.putAll($executorIdToItemIndexListMap);
        } catch (Exception e) {
            log.warn("deserialize " + taskId + "'s shards from necessary failed, will try to get shards from sharding/content", e);
            executorIdToItemIndexListMap.putAll(getExecutorIdToItemIndexListMap(taskId));
        }
        return executorIdToItemIndexListMap;
    }

}
